import akshare as ak
import pandas as pd
import os

# 确保安装了 fastparquet
# pip install fastparquet

# 1. 获取所有概念板块列表
print("正在获取概念板块列表...")
concept_list = ak.stock_board_concept_name_em()
print(f"获取到 {len(concept_list)} 个概念板块。")

# 2. 遍历所有概念板块，获取成分股
all_concept_stocks = []
for concept_name in concept_list['板块名称']:
    try:
        print(f"正在获取 {concept_name} 的成分股...")
        stocks = ak.stock_board_concept_cons_em(concept_name)
        stocks['概念板块'] = concept_name  # 添加所属概念列
        all_concept_stocks.append(stocks)
        print(f"成功获取 {concept_name} 成分股。")
    except Exception as e:
        print(f"获取 {concept_name} 失败: {e}")

# 3. 合并所有数据
if not all_concept_stocks:
    print("没有获取到任何股票数据，程序退出。")
    exit()

final_df = pd.concat(all_concept_stocks, ignore_index=True)

# 4. 重命名列
# 首先查看原始列名，以便准确映射
print("\n原始 DataFrame 列名：")
print(final_df.columns)

# 根据 akshare 返回的实际列名进行映射
# 请注意：akshare接口返回的列名可能会有变动，请根据实际情况调整下面的映射
# 以下是一个根据常见返回值的猜测映射，你可能需要根据实际运行结果进行调整
column_mapping = {
    '代码': 'code',
    '名称': 'name',
    '最新价': 'close',  # 或者 'current_price'
    '涨跌幅': 'price_change_ratio',
    '涨跌额': 'price_change_amount',
    '成交量': 'volume',
    '成交额': 'amount',
    '振幅': 'amplitude',
    '最高': 'high',
    '最低': 'low',
    '今开': 'open',
    '昨收': 'last_close', # 昨收通常是 previous_close 或 last_close
    '换手率': 'turnover_rate',
    '市盈率-动态': 'pe_ratio_dynamic',
    '市净率': 'pb_ratio',
    '概念板块': 'industry' # 映射到你需要的 'industry'
}

# 过滤掉不在 DataFrame 中的列，避免 KeyError
actual_column_mapping = {k: v for k, v in column_mapping.items() if k in final_df.columns}

# 执行重命名
final_df.rename(columns=actual_column_mapping, inplace=True)

# 确保包含你需要的英文列
required_columns = ["code", "name", "open", "high", "low", "close", "volume", "amount", "industry"]
# 检查是否所有必需的列都存在，如果不存在则填充为 NaN 或者进行其他处理
for col in required_columns:
    if col not in final_df.columns:
        print(f"警告: 缺少必需的列 '{col}'，已添加为空列。")
        final_df[col] = pd.NA # 或者 None, np.nan 等

# 仅保留你需要的列，并调整顺序
# 你可以根据实际情况添加 'capital' 和 'turnover'，如果 akshare 接口不直接提供，可能需要计算或从其他接口获取
# 这里我假设 'turnover' 对应 '换手率' 重命名后的 'turnover_rate'
# 'capital' (总股本/流通市值) 一般需要额外数据来计算，这里暂时无法直接从现有数据获得。
# 如果你想保留更多字段，请修改此列表
output_columns = [
    "code", "name", "open", "high", "low", "close", "volume", "amount", 
    "turnover_rate", # 将换手率保留下来，你可以将其映射到 'turnover' 如果觉得合适
    "industry"
]

# 筛选并重新排序列
final_df = final_df[output_columns]

# 如果需要将 'turnover_rate' 重新命名为 'turnover'
if 'turnover_rate' in final_df.columns:
    final_df.rename(columns={'turnover_rate': 'turnover'}, inplace=True)

# 打印最终的 DataFrame 信息
print("\n最终 DataFrame 预览：")
print(final_df.head())
print("\n最终 DataFrame 列名和数据类型：")
print(final_df.info())

# 5. 保存为 Parquet 文件
output_filename = "all_concept_stocks.parquet"
try:
    final_df.to_parquet(output_filename, engine='fastparquet', index=False)
    print(f"\n数据已成功保存到 {output_filename}")
except Exception as e:
    print(f"保存 Parquet 文件失败: {e}")